import os
import json
import torch
import torch.nn as nn
import numpy as np
import core.pyplot_utils as utils
from torch.utils.data import DataLoader

from tensorboardX import SummaryWriter
from pandas.plotting import register_matplotlib_converters

from core.data_processor import readData
from core.model import mytrainset, RNN,GRU,LSTM




def run():
    #数据装配
    configs = json.load(open('config.json', 'r'))
    path = os.path.join('data', configs['data']['filename'])
    stockname=configs['data']['stockname']
    input_dim = configs['data']['input_dim']
    hidden_dim= configs['data']['hidden_dim']
    num_layers= configs['data']['num_layers']
    output_dim=configs['data']['output_dim']
    lr = configs['data']['lr']
    epoch = configs['data']['epoch']
    batch_size = configs['data']['batch_size']
    train_end = configs['data']['train_end']

    modeltype=configs['model']['type']
    columns = 'close'

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


    register_matplotlib_converters()

    # 获取训练数据、原始数据、索引等信息
    df, df_all, df_index = readData(path, columns, input_dim,True, False, train_end)

    # 对数据进行预处理，规范化及转换为Tensor
    df_numpy = np.array(df)

    df_numpy_mean = np.mean(df_numpy)
    df_numpy_std = np.std(df_numpy)

    df_numpy = (df_numpy - df_numpy_mean) / df_numpy_std
    df_tensor = torch.Tensor(df_numpy)

    trainset = mytrainset(df_tensor)
    trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=False)

    writer = SummaryWriter(log_dir='logs')

    #rnn = LSTM(n, 64, 1, 1).to(device)
    if modeltype == 'GRU':
     rnn = GRU(input_dim, hidden_dim, num_layers, output_dim).to(device)
     print('GRU--------------------------------------------')

    if modeltype == 'LSTM':
     rnn = LSTM(input_dim, hidden_dim, num_layers, output_dim).to(device)
     print('LSTM--------------------------------------------')

    optimizer = torch.optim.Adam(rnn.parameters(), lr)
    loss_func = nn.MSELoss()

    for step in range(epoch):
        for tx, ty in trainloader:
            tx = tx.to(device)
            ty = ty.to(device)
            # 在第1个维度上添加一个维度为1的维度，形状变为[batch,seq_len,input_size]
            output = rnn(torch.unsqueeze(tx, dim=1)).to(device)
            loss = loss_func(torch.squeeze(output), ty)
            print("Epoch ", step, "MSE: ", loss.item())

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        writer.add_scalar('loss', loss, step)

    generate_data_train = []
    generate_data_test = []

    test_index = len(df_all) + train_end

    df_all_normal = (df_all - df_numpy_mean) / df_numpy_std
    df_all_normal_tensor = torch.Tensor(df_all_normal)
    for i in range(input_dim, len(df_all)):
        x = df_all_normal_tensor[i - input_dim:i].to(device)
        # rnn的输入必须是3维，故需添加两个1维的维度，最后成为[1,1,input_size]
        x = torch.unsqueeze(torch.unsqueeze(x, dim=0), dim=0)

        y = rnn(x).to(device)
        if i < test_index:
            generate_data_train.append(torch.squeeze(y).detach().cpu().numpy() * df_numpy_std + df_numpy_mean)
        else:
            generate_data_test.append(torch.squeeze(y).detach().cpu().numpy() * df_numpy_std + df_numpy_mean)

    utils.plot_results(modeltype, stockname, df_index, df_all, generate_data_train, input_dim, generate_data_test, train_end)

    utils.plot_results_custom(modeltype, stockname, df_index, df_all, generate_data_test, -500, -1)




if __name__ == '__main__':
    run()

